# https://blog.csdn.net/m0_37383484/article/details/133387751
from ultralytics import YOLO
import cv2
import base64

def drawImage(img,regionList):
    for r in regionList:
        x,y,w,h = r['bbox']
        type = r['type']
        text = type
        cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
    imgStr = cv2.imencode('.jpg', img)[1].tostring()
    imgBase64 = base64.b64encode(imgStr).decode()
    return imgBase64

def getNodeInfo(img):
    model_path = '/data/zwxu/layoutparse/model/general6-8n.pt'  # 权重路径
    model = YOLO(model_path)
    regionList = []
    result = model(img, save=False, conf=0.5, save_crop=False, line_width=2)
    for r in result:
        names = r.names
        boxes = r.boxes
        for box in boxes:
            # Bounding Box
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            w, h = x2 - x1, y2 - y1
            # Confidence
            #conf = math.ceil((box.conf[0] * 100)) / 100
            # Class Name
            cls = int(box.cls[0])
            box = x1,y1,w,h
            region = {'type':names[cls],'bbox': box}
            regionList.append(region)

    height,width,_ = img.shape
    imgBase64 = drawImage(img,regionList)
    nodeInfo = {'height':height,'width':width,'node':regionList,'base64':imgBase64}
    return nodeInfo     

#image_path = './0722.jpg'  # 待预测图片路径
#img = cv2.imread(image_path)
#nodeInfo = getNodeInfo(img)
#print(nodeInfo)
